#https://zhuanlan.zhihu.com/p/51400008
#https://zhuanlan.zhihu.com/p/61628249

import heapq
def dijkstra_heap(graph, vertex_num, s):
    h = []
    dist = {}
    flag = [False] * vertex_num
    for i in range(vertex_num):
        dist[i] = 127
    # 将设置的起始端点加进来
    dist[s] = 0
    heapq.heappush(h, (dist[s], s))
    while len(h) != 0:
        mini_dist = heapq.nsmallest(1, h)[0]
        heapq.heappop(h)
        tmp = mini_dist[1]
        # if主要是为了减少复杂度，已经做过松弛的节点不需要再做松弛了
        if not flag[tmp]:
            flag[tmp] = True
            for j in range(len(graph[tmp])):
                if dist[j] <= dist[tmp] + graph[tmp][j]:
                    continue
                dist[j] = dist[tmp] + graph[tmp][j]
                heapq.heappush(h, (dist[j], j))
    return dist
    
def dijkstra(matrix_distance, source_node):
    inf = float('inf')
    # init the source node distance to others
    dis = matrix_distance[source_node]
    node_nums = len(dis)
    
    flag = [0 for i in range(node_nums)]
    flag[source_node] = 1
    
    for i in range(node_nums-1):
        min = inf
        #find the min node from the source node
        for j in range(node_nums):
            if flag[j] == 0 and dis[j] < min:
                min = dis[j]
                u = j
        flag[u] = 1
        #update the dis 
        for v in range(node_nums):
            if flag[v] == 0 and matrix_distance[u][v] < inf:
                if dis[v] > dis[u] + matrix_distance[u][v]:
                    dis[v] = dis[u] + matrix_distance[u][v]                    
    
    return dis



if __name__ == '__main__':
    graph = [[0, 1, 10, 999],
             [1, 0, 2,  999],
             [10, 2, 0, 3, ],
             [999, 999, 3, 0]]
    print(dijkstra_heap(graph, 4, 1))
    print(dijkstra(matrix_distance, 0))
    pass